"""
查询计划和执行树结构
"""

from typing import List, Dict, Any, Optional, Union
from dataclasses import dataclass, field
from enum import Enum
import json
import uuid

from .types import (
    OperatorType, JoinType, AggregateFunction, SortDirection,
    DataType, ExecutionContext, ExecutionStats
)

@dataclass
class PlanNode:
    """查询计划节点"""
    node_id: str
    operator_type: OperatorType # 操作符类型
    properties: Dict[str, Any] = field(default_factory=dict) # 属性
    children: List['PlanNode'] = field(default_factory=list) # 子节点
    parent: Optional['PlanNode'] = None # 父节点
    estimated_cost: float = 0.0 # 估计成本
    estimated_rows: int = 0 # 估计行数
    estimated_selectivity: float = 1.0 # 估计选择性
    
    def __post_init__(self):
        if not self.node_id:
            self.node_id = str(uuid.uuid4())
        
        # 设置父子关系
        for child in self.children:
            child.parent = self
    
    def add_child(self, child: 'PlanNode'):
        """添加子节点"""
        self.children.append(child)
        child.parent = self
    
    def remove_child(self, child: 'PlanNode'):
        """移除子节点"""
        if child in self.children:
            self.children.remove(child)
            child.parent = None
    
    def get_property(self, key: str, default=None):
        """获取属性值"""
        return self.properties.get(key, default)
    
    def set_property(self, key: str, value: Any):
        """设置属性值"""
        self.properties[key] = value
    
    def is_leaf(self) -> bool:
        """是否为叶子节点"""
        return len(self.children) == 0
    
    def is_root(self) -> bool:
        """是否为根节点"""
        return self.parent is None
    
    def depth(self) -> int:
        """节点深度"""
        if self.is_root():
            return 0
        return self.parent.depth() + 1
    
    def subtree_size(self) -> int:
        """子树大小"""
        size = 1
        for child in self.children:
            size += child.subtree_size()
        return size
    
    def to_dict(self) -> Dict[str, Any]:
        """转换为字典"""
        return {
            'node_id': self.node_id,
            'operator_type': self.operator_type.value,
            'properties': self.properties,
            'children': [child.to_dict() for child in self.children],
            'estimated_cost': self.estimated_cost,
            'estimated_rows': self.estimated_rows,
            'estimated_selectivity': self.estimated_selectivity
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'PlanNode':
        """从字典创建节点"""
        node = cls(
            node_id=data['node_id'],
            operator_type=OperatorType(data['operator_type']),
            properties=data.get('properties', {}),
            estimated_cost=data.get('estimated_cost', 0.0),
            estimated_rows=data.get('estimated_rows', 0),
            estimated_selectivity=data.get('estimated_selectivity', 1.0)
        )
        
        # 递归创建子节点
        for child_data in data.get('children', []):
            child = cls.from_dict(child_data)
            node.add_child(child)
        
        return node

class QueryPlan:
    """查询计划"""
    
    def __init__(self, plan_id: str = None):
        self.plan_id = plan_id or str(uuid.uuid4()) # 计划ID
        self.root_node: Optional[PlanNode] = None
        self.metadata: Dict[str, Any] = {} # 元数据
        self.optimization_level: int = 1 # 优化级别
        self.creation_time = None # 创建时间
        self.estimated_total_cost: float = 0.0
        
    def set_root(self, root: PlanNode):
        """设置根节点"""
        self.root_node = root
        self._update_cost_estimates()
    
    def _update_cost_estimates(self):
        """更新成本估算"""
        if self.root_node:
            self.estimated_total_cost = self._calculate_subtree_cost(self.root_node)
    
    def _calculate_subtree_cost(self, node: PlanNode) -> float:
        """计算子树成本"""
        cost = node.estimated_cost
        for child in node.children:
            cost += self._calculate_subtree_cost(child)
        return cost
    
    def get_all_nodes(self) -> List[PlanNode]:
        """获取所有节点"""
        if not self.root_node:
            return []
        
        nodes = []
        self._collect_nodes(self.root_node, nodes)
        return nodes
    
    def _collect_nodes(self, node: PlanNode, nodes: List[PlanNode]):
        """递归收集节点"""
        nodes.append(node)
        for child in node.children:
            self._collect_nodes(child, nodes)
    
    def find_node_by_id(self, node_id: str) -> Optional[PlanNode]:
        """根据ID查找节点"""
        for node in self.get_all_nodes():
            if node.node_id == node_id:
                return node
        return None
    
    def find_nodes_by_type(self, operator_type: OperatorType) -> List[PlanNode]:
        """根据操作符类型查找节点"""
        return [node for node in self.get_all_nodes() 
                if node.operator_type == operator_type]
    
    def validate(self) -> List[str]:
        """验证查询计划"""
        errors = []
        
        if not self.root_node:
            errors.append("Query plan has no root node")
            return errors
        
        # 验证节点连接性
        nodes = self.get_all_nodes()
        for node in nodes:
            # 验证父子关系一致性
            for child in node.children:
                if child.parent != node:
                    errors.append(f"Inconsistent parent-child relationship for node {node.node_id}")
        
        # 验证操作符特定的约束
        for node in nodes:
            node_errors = self._validate_node(node)
            errors.extend(node_errors)
        
        return errors
    
    def _validate_node(self, node: PlanNode) -> List[str]:
        """验证单个节点"""
        errors = []
        
        # 根据操作符类型进行特定验证
        if node.operator_type in [OperatorType.HASH_JOIN, OperatorType.NESTED_LOOP_JOIN, 
                                 OperatorType.SORT_MERGE_JOIN]:
            if len(node.children) != 2:
                errors.append(f"Join node {node.node_id} must have exactly 2 children")
            
            join_type = node.get_property('join_type')
            if not join_type or not isinstance(join_type, JoinType):
                errors.append(f"Join node {node.node_id} must have valid join_type property")
        
        elif node.operator_type == OperatorType.TABLE_SCAN:
            table_name = node.get_property('table_name')
            if not table_name:
                errors.append(f"Table scan node {node.node_id} must have table_name property")
        
        elif node.operator_type == OperatorType.FILTER:
            predicate = node.get_property('predicate')
            if not predicate:
                errors.append(f"Filter node {node.node_id} must have predicate property")
        
        return errors
    
    def optimize(self, optimization_rules: List[str] = None):
        """优化查询计划"""
        if not self.root_node:
            return
        
        # 默认优化规则
        if optimization_rules is None:
            optimization_rules = [
                'push_down_predicates',# 谓词下推优化
                'push_down_projections', # 投影下推优化
                'join_reordering',# 连接重排序优化
                'index_selection'# 索引选择优化
            ]
        
        for rule in optimization_rules:
            self._apply_optimization_rule(rule)
        
        self._update_cost_estimates()
    
    def _apply_optimization_rule(self, rule: str):
        """应用优化规则"""
        if rule == 'push_down_predicates':
            self._push_down_predicates()
        elif rule == 'push_down_projections':
            self._push_down_projections()
        elif rule == 'join_reordering':
            self._reorder_joins()
        elif rule == 'index_selection':
            self._select_indexes()
    
    def _push_down_predicates(self):
        """谓词下推优化"""
        # 简化的谓词下推实现
        pass
    
    def _push_down_projections(self):
        """投影下推优化"""
        # 简化的投影下推实现
        pass
    
    def _reorder_joins(self):
        """连接重排序优化"""
        # 简化的连接重排序实现
        pass
    
    def _select_indexes(self):
        """索引选择优化"""
        # 简化的索引选择实现
        pass
    
    def to_json(self) -> str:
        """转换为JSON字符串"""
        plan_dict = {
            'plan_id': self.plan_id,
            'root_node': self.root_node.to_dict() if self.root_node else None,
            'metadata': self.metadata,
            'optimization_level': self.optimization_level,
            'estimated_total_cost': self.estimated_total_cost
        }
        return json.dumps(plan_dict, indent=2)
    
    @classmethod
    def from_json(cls, json_str: str) -> 'QueryPlan':
        """从JSON字符串创建查询计划"""
        data = json.loads(json_str)
        
        plan = cls(data['plan_id'])
        plan.metadata = data.get('metadata', {})
        plan.optimization_level = data.get('optimization_level', 1)
        plan.estimated_total_cost = data.get('estimated_total_cost', 0.0)
        
        if data.get('root_node'):
            plan.root_node = PlanNode.from_dict(data['root_node'])
        
        return plan
    
    def explain(self, format_type: str = 'text') -> str:
        """生成执行计划说明"""
        if not self.root_node:
            return "Empty query plan"
        
        if format_type == 'text':
            return self._explain_text()
        elif format_type == 'json':
            return self.to_json()
        else:
            raise ValueError(f"Unsupported format type: {format_type}")
    
    def _explain_text(self) -> str:
        """生成文本格式的执行计划说明"""
        lines = []
        lines.append(f"Query Plan (ID: {self.plan_id})")
        lines.append(f"Estimated Total Cost: {self.estimated_total_cost:.2f}")
        lines.append("-" * 50)
        
        if self.root_node:
            self._explain_node_text(self.root_node, lines, 0)
        
        return "\n".join(lines)
    
    def _explain_node_text(self, node: PlanNode, lines: List[str], indent: int):
        """生成节点的文本说明"""
        prefix = "  " * indent + "└─ " if indent > 0 else ""
        
        # 基本信息
        line = f"{prefix}{node.operator_type.value.upper()}"
        
        # 添加重要属性
        important_props = []
        if node.operator_type == OperatorType.TABLE_SCAN:
            table_name = node.get_property('table_name')
            if table_name:
                important_props.append(f"table={table_name}")
        
        elif node.operator_type in [OperatorType.HASH_JOIN, OperatorType.NESTED_LOOP_JOIN]:
            join_type = node.get_property('join_type')
            if join_type:
                    important_props.append(f"type={join_type}")
            
            join_condition = node.get_property('join_condition')
            if join_condition:
                important_props.append(f"on={join_condition}")
        
        elif node.operator_type == OperatorType.FILTER:
            predicate = node.get_property('predicate')
            if predicate:
                important_props.append(f"where={predicate}")
        
        if important_props:
            line += f" ({', '.join(important_props)})"
        
        # 成本信息
        line += f" [cost={node.estimated_cost:.2f}, rows={node.estimated_rows}]"
        
        lines.append(line)
        
        # 递归处理子节点
        for child in node.children:
            self._explain_node_text(child, lines, indent + 1)

@dataclass
class ExecutionNode:
    """执行节点，包含运行时信息"""
    plan_node: PlanNode # 计划节点
    executor_id: str
    context: ExecutionContext # 执行上下文
    runtime_stats: ExecutionStats = field(default_factory=ExecutionStats) # 运行时统计
    children: List['ExecutionNode'] = field(default_factory=list) # 子节点
    parent: Optional['ExecutionNode'] = None
    
    def __post_init__(self):
        # 设置父子关系
        for child in self.children:
            child.parent = self
    
    def add_child(self, child: 'ExecutionNode'):
        """添加子执行节点"""
        self.children.append(child)
        child.parent = self
    
    def update_stats(self, stats: ExecutionStats):
        """更新运行时统计"""
        self.runtime_stats = stats
    
    def get_total_stats(self) -> ExecutionStats:
        """获取包含子节点的总统计"""
        total = ExecutionStats(
            rows_processed=self.runtime_stats.rows_processed,
            rows_returned=self.runtime_stats.rows_returned,
            execution_time_ms=self.runtime_stats.execution_time_ms,
            memory_used_bytes=self.runtime_stats.memory_used_bytes,
            io_operations=self.runtime_stats.io_operations,
            cpu_time_ms=self.runtime_stats.cpu_time_ms
        )
        
        for child in self.children:
            child_stats = child.get_total_stats()
            total.rows_processed += child_stats.rows_processed
            total.rows_returned += child_stats.rows_returned
            total.execution_time_ms += child_stats.execution_time_ms
            total.memory_used_bytes += child_stats.memory_used_bytes
            total.io_operations += child_stats.io_operations
            total.cpu_time_ms += child_stats.cpu_time_ms
        
        return total

class PlanBuilder:
    """查询计划构建器"""
    
    def __init__(self):
        self.current_plan = QueryPlan()
    
    def scan_table(self, table_name: str, alias: str = None) -> PlanNode:
        """创建表扫描节点"""
        node = PlanNode(
            node_id=str(uuid.uuid4()),
            operator_type=OperatorType.TABLE_SCAN
        )
        node.set_property('table_name', table_name)
        if alias:
            node.set_property('alias', alias)
        
        return node
    
    def filter(self, child: PlanNode, predicate: str) -> PlanNode:
        """创建过滤节点"""
        node = PlanNode(
            node_id=str(uuid.uuid4()),
            operator_type=OperatorType.FILTER,
            children=[child]
        )
        node.set_property('predicate', predicate)
        return node
    
    def project(self, child: PlanNode, columns: List[str]) -> PlanNode:
        """创建投影节点"""
        node = PlanNode(
            node_id=str(uuid.uuid4()),
            operator_type=OperatorType.PROJECT,
            children=[child]
        )
        node.set_property('columns', columns)
        return node
    
    def join(self, left: PlanNode, right: PlanNode, 
            join_type: JoinType, condition: str) -> PlanNode:
        """创建连接节点"""
        node = PlanNode(
            node_id=str(uuid.uuid4()),
            operator_type=OperatorType.HASH_JOIN,
            children=[left, right]
        )
        node.set_property('join_type', join_type)
        node.set_property('join_condition', condition)
        return node
    
    def aggregate(self, child: PlanNode, group_by: List[str], 
                 aggregates: List[Dict[str, Any]]) -> PlanNode:
        """创建聚合节点"""
        node = PlanNode(
            node_id=str(uuid.uuid4()),
            operator_type=OperatorType.HASH_AGGREGATE,
            children=[child]
        )
        node.set_property('group_by', group_by)
        node.set_property('aggregates', aggregates)
        return node
    
    def sort(self, child: PlanNode, order_by: List[Dict[str, Any]]) -> PlanNode:
        """创建排序节点"""
        node = PlanNode(
            node_id=str(uuid.uuid4()),
            operator_type=OperatorType.SORT,
            children=[child]
        )
        node.set_property('order_by', order_by)
        return node
    
    def limit(self, child: PlanNode, limit: int, offset: int = 0) -> PlanNode:
        """创建限制节点"""
        node = PlanNode(
            node_id=str(uuid.uuid4()),
            operator_type=OperatorType.LIMIT,
            children=[child]
        )
        node.set_property('limit', limit)
        node.set_property('offset', offset)
        return node
    
    def build(self, root: PlanNode) -> QueryPlan:
        """构建查询计划"""
        self.current_plan.set_root(root)
        return self.current_plan

